#! ~/anaconda3/bin/python
# -*- coding: utf-8 -*-
# @Time   : 2019-11-15
# @Author : chujie
# @Email  : zjchenchujie@126.com
# @Brief  : find the annotations that unmatch image.
"""
1. 遮挡指定区域
2. 若遮挡区域存在目标，则不遮挡目标区域
"""

import xml.etree.ElementTree as ET
import os
import cv2
import math
import argparse
import shutil


def count_files(img_dir):
    """
    count all the file in a directory, recursively
    :param dir:
    :return:
    """
    # return sum([len(files) for root, dirs, files in os.walk(dir)])
    return len(os.listdir(img_dir))


def process(img_file, xml_file, processed_img_save_dir):
    """
    main process function.
    :param img_file: image input
    :param xml_file: xml file
    :param processed_img_save_dir: target saving directory
    :return: None
    """
    img_input = cv2.imread(img_file)
    img_output = cv2.imread(img_file)
    # rectangle to be covered [xmin, ymin, xmax, ymax]
    # ignore_area=[[90,280,165,495],[500,500,704,576]]
    ignore_area = [[500, 500, 704, 576]]
    for i in range(len(ignore_area)):
        img_output[ignore_area[i][1]:ignore_area[i][3], ignore_area[i][0]:ignore_area[i][2]] = 0
    tree = ET.parse(xml_file)
    root = tree.getroot()
    height, width, _ = img_input.shape
        
    for obj in root.iter('object'):
        xml_box = obj.find('bndbox')
        left = int(xml_box.find('xmin').text)
        top = int(xml_box.find('ymin').text)
        right = int(xml_box.find('xmax').text)
        bottom = int(xml_box.find('ymax').text)
        target_width = right - left
        target_height = bottom - top

        # 扩展像素
        e_limit = 4
        ext_x = math.floor(0.1 * target_width)
        if ext_x < e_limit:
            ext_x = e_limit
        elif ext_x > 2 * e_limit:
            ext_x = 2 * e_limit
        ext_y = math.floor(0.1 * target_height)
        if ext_y < e_limit:
            ext_y = e_limit
        elif ext_y > 2 * e_limit:
            ext_y = 2 * e_limit
            
        if left - ext_x > 0:
            left = left - ext_x
        else:
            left = 0
        
        if top - ext_y > 0:
            top = top - ext_y
        else:
            top = 0
            
        if right + ext_x < width:
            right = right + ext_x
        else:
            right = width
        
        if bottom + ext_y < height:
            bottom = bottom + ext_y
        else:
            bottom = height
        img_output[top:bottom, left:right] = img_input[top:bottom, left:right]
    img_name = img_file.split('/')[-1]
    save_path = os.path.join(processed_img_save_dir, img_name)
    cv2.imwrite(save_path, img_output)


def main():
    """
    main
    :return:
    """
    root_dir = os.path.join(os.getenv("HOME"), "PycharmProjects/data_tool/data/wider_face_VOC")
    parser = argparse.ArgumentParser(description="define all the file paths")
    parser.add_argument("--img_dir", type=str, help="image directory")
    parser.add_argument("--xml_dir", type=str, help="annotation directory")
    parser.add_argument("--save_dir", type=str, help="save directory")
    args = parser.parse_args()
    img_dir = args.img_dir
    if not img_dir:
        print("argument  \"img_dir\" missing, use default \"{}\" as default."
              .format(os.path.join(root_dir, "JPEGImages")))
        if not os.path.exists(os.path.join(root_dir, "JPEGImages")):
            print("image directory \"{}\" not found. ".format(os.path.join(root_dir, "JPEGImages")))
            exit()
        img_dir = os.path.join(root_dir, "JPEGImages")
    xml_dir = args.xml_dir
    if not xml_dir:
        print("argument  \"xml_dir\" missing, use default \"{}\" as default."
              .format(os.path.join(root_dir, "Annotations")))
        if not os.path.exists(os.path.join(root_dir, "Annotations")):
            print("image directory \"{}\" not found. ".format(os.path.join(root_dir, "Annotations")))
            exit()
        xml_dir = os.path.join(root_dir, "Annotations")
    save_dir = args.save_dir
    if not save_dir:
        print("argument  \"save_dir\" missing, use default \"{}\" as default."
              .format(os.path.join(root_dir, "save_dir")))
        if os.path.exists(os.path.join(root_dir, "save_dir")):
            shutil.rmtree(os.path.join(root_dir, "save_dir"))
            os.mkdir(os.path.join(root_dir, "save_dir"))
        else:
            os.mkdir(os.path.join(root_dir, "save_dir"))
        save_dir = os.path.join(root_dir, "save_dir")
    
    file_num = count_files(img_dir)
    processed_num = 0
    files = os.listdir(img_dir)
    for file in files:
        img_name = file.split('.')[0]
        img_file = os.path.join(img_dir, file)
        xml_file = os.path.join(xml_dir, '{}.xml'.format(img_name))
        process(img_file, xml_file, save_dir)
        processed_num = processed_num + 1
        if file_num % 100 == 0:
            print('Processed {} \\ {} '.format(processed_num, file_num))
    print('completed all images :'.format(processed_num))


if __name__ == '__main__':
    main()
